/*
 *    Copyright 2022 The DSMS Authors.
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *        http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */
package com.dsms.common.taskmanager.job;

import com.alibaba.fastjson2.JSON;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.dsms.common.constant.TaskExclusiveEnum;
import com.dsms.common.constant.TaskStatusEnum;
import com.dsms.common.taskmanager.TaskContext;
import com.dsms.common.taskmanager.model.Task;
import com.dsms.common.taskmanager.service.ITaskService;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.ListMultimap;
import javax.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import org.springframework.util.ObjectUtils;
import org.springframework.util.StringUtils;

import java.time.LocalDateTime;
import java.util.List;
import java.util.Objects;

@Slf4j
@Component
public class TaskJob {
    private final ITaskService taskService;
    private final TaskContext taskContext;
    private final ThreadPoolTaskExecutor dsmsExecutor;

    @Autowired
    public TaskJob(ITaskService taskService, TaskContext taskContext, ThreadPoolTaskExecutor dsmsExecutor){
        this.taskService = taskService;
        this.taskContext = taskContext;
        this.dsmsExecutor = dsmsExecutor;
    }

    @Scheduled(cron = "0/2 * * * * ? ")
    public void searchTask() {
        log.debug("TaskJob execute ....");
        //if cluster not bind, skip task execute
        if (!taskService.validateBindCluster()) {
            return;
        }
        //get the number of idle threads in the thread pool
        int idleThreadNum = dsmsExecutor.getThreadPoolExecutor().getQueue().remainingCapacity() - dsmsExecutor.getActiveCount();
        if (idleThreadNum == 0) {
            return;
        }
        //get task in queue or rollback status,the task number we get is the same as the number of idle threads
        LambdaQueryWrapper<Task> queryWrapper = new LambdaQueryWrapper<Task>().eq(Task::getTaskStatus, TaskStatusEnum.QUEUE.getStatus())
                .or().eq(Task::getTaskStatus, TaskStatusEnum.ROLLBACK.getStatus()).last("limit " + idleThreadNum);
        List<Task> tasks = taskService.list(queryWrapper);
        if (ObjectUtils.isEmpty(tasks)) {
            return;
        }

        //preExecuteTaskList contains all non-exclusive tasks, and only one exclusive task of the same type.
        ListMultimap<String, Task> preExecuteTaskList = ArrayListMultimap.create();
        for (Task task : tasks) {
            String taskType = task.getTaskType();
            boolean isExclusive = Objects.equals(TaskExclusiveEnum.EXCLUSIVE.getCode(), task.getIsExclusive());
            if (!isExclusive || !preExecuteTaskList.containsKey(task.getTaskType())) {
                preExecuteTaskList.put(taskType, task);
            }
        }

        for (Task task : preExecuteTaskList.values()) {
            //if it is an exclusive task, query whether there is a task of the same type being executed, and if so, skip this task
            if (Objects.equals(task.getIsExclusive(), TaskExclusiveEnum.EXCLUSIVE.getCode())) {
                LambdaQueryWrapper<Task> executingQueryWrapper = new LambdaQueryWrapper<Task>().eq(Task::getTaskStatus, TaskStatusEnum.EXECUTING.getStatus())
                        .eq(Task::getTaskType, task.getTaskType());
                List<Task> sameTypeTasksInExecuting = taskService.list(executingQueryWrapper);
                if (!sameTypeTasksInExecuting.isEmpty()) {
                    continue;
                }
            }
            dsmsExecutor.submit(() -> {
                try {
                    //roll back failed task
                    if (Objects.equals(task.getTaskStatus(), TaskStatusEnum.ROLLBACK.getStatus())) {
                        taskContext.rollback(task);
                    } else {
                        //execute queue task
                        //Update the task status to Executing, prevent repeated execution
                        task.setTaskStatus(TaskStatusEnum.EXECUTING.getStatus());
                        taskService.updateById(task);
                        taskContext.execute(task);
                    }
                } catch (Throwable t) {
                    log.error("execute task fail,fail message:{},task info:{}", t.getMessage(), JSON.toJSONString(task), t);
                    task.setTaskStatus(TaskStatusEnum.FAIL.getStatus());
                    task.setTaskErrorMessage(StringUtils.hasText(task.getTaskErrorMessage()) ? task.getTaskErrorMessage() : t.getMessage());
                } finally {
                    try {
                        if (TaskStatusEnum.FAIL.getStatus() == task.getTaskStatus()
                                || TaskStatusEnum.FINISH.getStatus() == task.getTaskStatus()
                                || TaskStatusEnum.ROLLBACK_SUCCESS.getStatus() == task.getTaskStatus()) {
                            task.setTaskEndTime(LocalDateTime.now());
                        }
                        taskService.updateById(task);
                        taskService.updateFrontTaskInfo(task);
                    } catch (Throwable t) {
                        log.error("update task info fail,fail message:{},task info:{}", t.getMessage(), JSON.toJSONString(task), t);
                    }
                }
            });
        }
    }

}
